package cn.ai.shoes.client.service.impl;

import cn.ai.shoes.client.api.ComfyuiApi;
import cn.ai.shoes.client.config.ComfyuiConfig;
import cn.ai.shoes.client.entity.ComfyuiUploadLog;
import cn.ai.shoes.client.entity.OriginalRefModel;
import cn.ai.shoes.client.entity.OssUploadLog;
import cn.ai.shoes.client.entity.CurrentServerEntity;
import cn.ai.shoes.client.mapper.ComyuiUploadMapper;
import cn.ai.shoes.client.service.ComfyuiUploadService;
import cn.ai.shoes.client.service.OssUploadLogService;
import cn.ai.shoes.client.service.OssUploadService;
import cn.ai.shoes.common.exception.ServiceException;
import cn.ai.shoes.common.exception.constant.GlobalErrorCodeConstants;
import cn.ai.shoes.common.pojo.entity.CommonResult;
import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.util.ObjectUtil;
import com.alibaba.fastjson2.JSON;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import jakarta.annotation.Resource;
import jodd.util.StringUtil;
import lombok.extern.slf4j.Slf4j;
import okhttp3.MediaType;
import okhttp3.MultipartBody;
import okhttp3.OkHttpClient;
import okhttp3.RequestBody;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import retrofit2.Call;
import retrofit2.Response;
import retrofit2.Retrofit;
import retrofit2.converter.jackson.JacksonConverterFactory;

import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.HashMap;

@Service
@Slf4j
public class ComfyuiUploadServiceImpl extends ServiceImpl<ComyuiUploadMapper, ComfyuiUploadLog> implements ComfyuiUploadService {
    // 允许上传的格式 图片形式
    private static final String[] IMAGE_TYPE = new String[]{".bmp", ".jpg", ".jpeg", ".png"};
    public static final String INPUT = "input";
    public static final String CLIPSPACE = "clipspace";

    @Resource
    private CurrentServerEntity currentServerEntity;
    @Resource
    private OkHttpClient okHttpClient;
    @Resource
    private OssUploadService ossUploadService;
    @Resource
    private OssUploadLogService ossUploadLogService;

    /**
     * 上传图片到comfyui
     * @param ossUploadLog     图片日志信息
     * @param currentServerUrl 当前服务器地址
     * @return 上传结果
     */
    @Override
    @Transactional(rollbackFor = Exception.class)
    public ComfyuiUploadLog uploadImage(OssUploadLog ossUploadLog, String currentServerUrl,OssUploadLog originalOssUploadLog) {
        ComfyuiUploadLog comfyuiUploadLog = new ComfyuiUploadLog();
        //判断originalOssUploadLog是否为空
        if (ObjectUtil.isEmpty(originalOssUploadLog)) {
           comfyuiUploadLog = handUploadImage(ossUploadLog, currentServerUrl);
        }else {
            comfyuiUploadLog = handUploadMask(ossUploadLog, currentServerUrl,originalOssUploadLog);
        }
        return comfyuiUploadLog;
    }

    @NotNull
    private ComfyuiUploadLog handUploadImage(OssUploadLog ossUploadLog, String currentServerUrl) {
        ComfyuiUploadLog comfyuiUploadLog = new ComfyuiUploadLog();
        String cfUploadIds = ossUploadLog.getCfUploadIds();
        Long comfyuiUploadLogId = null;
        // 获取阿里云上的图片
        try (InputStream stream = ossUploadService.getFileInputStream(ossUploadLog.getOssPath())) {
            CommonResult<HashMap> response = null;
            // 转换为RequestBody
            RequestBody requestFile = convertBody(stream);
            // 构建MultipartBody.Part
            MultipartBody.Part filePart = MultipartBody.Part.createFormData("image", getNewFileName(ossUploadLog.getOssName()), requestFile);
            // 调用ComfyUI上传图片
            Call<HashMap> call = this.getComfyuiApi(currentServerUrl).uploadImage(filePart);
            // 拿到返回结果
            response = callComfyui(filePart, call);
            if (ObjectUtil.isNotEmpty(response) && response.getCode() == 200) {
                //成功
                HashMap data = response.getData();
                //将data的值按照key赋值给ComfyuiUploadLog
                comfyuiUploadLog = generateComfyuiUploadLog(currentServerUrl, data);
                //将comfyuiUploadLog保存到数据库
                ossUploadLog.setCfUploadIds(StringUtil.isNotEmpty(cfUploadIds) ? cfUploadIds + "," + (comfyuiUploadLogId != null ? comfyuiUploadLogId.toString() : "") : (comfyuiUploadLogId != null ? comfyuiUploadLogId.toString() : ""));
                return comfyuiUploadLog;
            }else{
                //失败
                log.error("上传图片到comfyui失败", response);
                comfyuiUploadLog.setIsSuccess(false);
                comfyuiUploadLog.setError(response.getMsg());
                throw new ServiceException(GlobalErrorCodeConstants.IMAGE_UPLOAD_ERROR);
            }
        } catch (Exception e) {
            //失败
            log.error("上传图片到comfyui失败", e);
            comfyuiUploadLog.setIsSuccess(false);
            comfyuiUploadLog.setError(e.getMessage());
            throw new ServiceException(GlobalErrorCodeConstants.IMAGE_UPLOAD_ERROR);
        } finally {
            try {
                save(comfyuiUploadLog);
                //将comfyuiUploadLog的id赋值给ossUploadLog
                comfyuiUploadLogId = comfyuiUploadLog.getId();
                //将comfyuiUploadLog保存到数据库
                ossUploadLog.setCfUploadIds(StringUtil.isNotEmpty(cfUploadIds) ? cfUploadIds + "," + (comfyuiUploadLogId != null ? comfyuiUploadLogId.toString() : "") : (comfyuiUploadLogId != null ? comfyuiUploadLogId.toString() : ""));
                //更新ossUploadLog
                ossUploadLogService.updateById(ossUploadLog);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
            return comfyuiUploadLog;
        }
    }

    /**
     * 上传图片到comfyui mask
     * @param ossUploadLog     图片日志信息
     * @param currentServerUrl 当前服务器地址
     * @return 上传结果
     */
    @Transactional(rollbackFor = Exception.class)
    public ComfyuiUploadLog handUploadMask(OssUploadLog ossUploadLog, String currentServerUrl,OssUploadLog originalOssUploadLog) {
        ComfyuiUploadLog comfyuiUploadLog = new ComfyuiUploadLog();
        ComfyuiUploadLog comfyuiOriginalUploadLog = handUploadImage(originalOssUploadLog, currentServerUrl);
        String cfUploadIds = ossUploadLog.getCfUploadIds();
        Long comfyuiUploadLogId = null;
        // 获取阿里云上的图片
        try (InputStream stream = ossUploadService.getFileInputStream(ossUploadLog.getOssPath())) {
            CommonResult<HashMap> response = null;
            // 转换为RequestBody
            RequestBody requestFile = convertBody(stream);
            // 构建MultipartBody.Part
            MultipartBody.Part filePart = MultipartBody.Part.createFormData("image", getNewFileName(ossUploadLog.getOssName()), requestFile);
            RequestBody type = RequestBody.create(INPUT, MediaType.parse("text/plain"));
            RequestBody subfolder = RequestBody.create(CLIPSPACE, MediaType.parse("text/plain"));
            RequestBody originalRef = RequestBody.create(
                    JSON.toJSONString(new OriginalRefModel(
                            null,
                            comfyuiOriginalUploadLog.getName(),
                            comfyuiOriginalUploadLog.getType())),
                    MediaType.parse("application/json")
            );
            // 调用ComfyUI上传图片
            Call<HashMap> call = getComfyuiApi(currentServerUrl).uploadMask(filePart, type, subfolder, originalRef);
            // 拿到返回结果
            response = callComfyui(filePart, call);
            if (ObjectUtil.isNotEmpty(response) && response.getCode() == 200) {
                //成功
                HashMap data = response.getData();
                //将data的值按照key赋值给ComfyuiUploadLog
                comfyuiUploadLog = generateComfyuiUploadLog(currentServerUrl, data);
                //将comfyuiUploadLog保存到数据库
                ossUploadLog.setCfUploadIds(StringUtil.isNotEmpty(cfUploadIds) ? cfUploadIds + "," + (comfyuiUploadLogId != null ? comfyuiUploadLogId.toString() : "") : (comfyuiUploadLogId != null ? comfyuiUploadLogId.toString() : ""));
                return comfyuiUploadLog;
            }else{
                //失败
                log.error("上传图片到comfyui失败", response);
                comfyuiUploadLog.setIsSuccess(false);
                comfyuiUploadLog.setError(response.getMsg());
                throw new ServiceException(GlobalErrorCodeConstants.IMAGE_UPLOAD_ERROR);
            }
        } catch (Exception e) {
            //失败
            log.error("上传图片到comfyui失败", e);
            comfyuiUploadLog.setIsSuccess(false);
            comfyuiUploadLog.setError(e.getMessage());
            throw new ServiceException(GlobalErrorCodeConstants.IMAGE_UPLOAD_ERROR);
        } finally {
            try {
                save(comfyuiUploadLog);
                //将comfyuiUploadLog的id赋值给ossUploadLog
                comfyuiUploadLogId = comfyuiUploadLog.getId();
                //将comfyuiUploadLog保存到数据库
                ossUploadLog.setCfUploadIds(StringUtil.isNotEmpty(cfUploadIds) ? cfUploadIds + "," + (comfyuiUploadLogId != null ? comfyuiUploadLogId.toString() : "") : (comfyuiUploadLogId != null ? comfyuiUploadLogId.toString() : ""));
                //更新ossUploadLog
                ossUploadLogService.updateById(ossUploadLog);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
            return comfyuiUploadLog;
        }
    }

    @NotNull
    private static ComfyuiUploadLog generateComfyuiUploadLog(String currentServerUrl, HashMap data) {
        ComfyuiUploadLog comfyuiUploadLog = BeanUtil.toBean(data, ComfyuiUploadLog.class);
        // 确保 id 字段为空，让数据库自动生成
        comfyuiUploadLog.setId(null);
        comfyuiUploadLog.setUrl(currentServerUrl);
        comfyuiUploadLog.setInUrl(true);
        comfyuiUploadLog.setIsSuccess(true);
        comfyuiUploadLog.setDeleted(0);
        return comfyuiUploadLog;
    }

    /**
     * 将 inputStream 转换为 RequestBody
     *
     * @param stream
     * @return
     */
    @NotNull
    private static RequestBody convertBody(InputStream stream) {
        RequestBody requestFile;
        try {
            requestFile = RequestBody.create(stream.readAllBytes(), MediaType.parse("image/*"));
        } catch (IOException e) {
            throw new ServiceException(GlobalErrorCodeConstants.IMAGE_UPLOAD_ERROR);
        }
        return requestFile;
    }

    /**
     * 调用ComfyUI上传图片
     *
     * @param filePart
     * @param call
     * @return
     * @throws Exception
     */
    @Nullable
    private CommonResult<HashMap> callComfyui(MultipartBody.Part filePart, Call<HashMap> call) throws Exception {
        Response<HashMap> response = null;
        try {
            response = call.execute();
        } catch (IOException e) {
            throw new ServiceException(GlobalErrorCodeConstants.IMAGE_UPLOAD_ERROR);
        }
        if (response.isSuccessful()) {
            return CommonResult.success(response.body());
        }
        throw new ServiceException(response.errorBody().string());
    }

    /**
     * 获取新的文件名
     *
     * @param fn
     * @return
     */
    @NotNull
    private static String getNewFileName(String fn) {
        String fileExtension = fn.substring(fn.lastIndexOf("."));
        if (!Arrays.asList(IMAGE_TYPE).contains(fileExtension.toLowerCase())) {
            throw new ServiceException(GlobalErrorCodeConstants.IMAGE_UPLOAD_ERROR);
        }
        String username = SecurityContextHolder.getContext().getAuthentication().getName();
        String randomSuffix = System.currentTimeMillis() + fileExtension;
        return randomSuffix;
    }

    // 每次动态获取最新服务器API实例
    private ComfyuiApi getComfyuiApi(String currentServerUrl) {
        //打印当前服务器地址
        log.info("[ComfyuiUploadServiceImpl.getComfyui]:当前服务器地址:"+ currentServerUrl);
        Retrofit retrofit = new Retrofit.Builder()
                .baseUrl(currentServerUrl)
                .client(okHttpClient)
                .addConverterFactory(JacksonConverterFactory.create())
                .build();
        return retrofit.create(ComfyuiApi.class);
    }

}
